from rest_framework_simplejwt.tokens import AccessToken, RefreshToken
from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
from django.utils.deprecation import MiddlewareMixin
import time
import logging
from rest_framework_simplejwt.settings import api_settings
from django.contrib.auth import get_user_model
from ..jwt.jwt_cache import (
    get_new_access_token,
    get_refresh_token,
    set_new_token,
    set_refresh_token,
)

logger = logging.getLogger("chat_app")


class TokenRefreshMiddleware(MiddlewareMixin):
    def process_response(self, request, response):
        if not 200 <= response.status_code < 300:
            return response

        auth_header = request.META.get("HTTP_AUTHORIZATION", "")
        if not auth_header or not auth_header.lower().startswith("bearer "):
            return response

        token_str = auth_header.split()[1]

        try:
            # 解析并验证基本有效性
            token = AccessToken(token_str)

            # 检查是否过期
            token.check_exp()  # 过期会抛InvalidToken

            # 计算剩余时间（秒）
            exp_time = token.payload["exp"]
            current_time = time.time()
            remaining = exp_time - current_time
            # 如果剩余时间小于15分钟则刷新
            if remaining < 900:
                refresh_token = get_refresh_token(token_str)
                new_access_token = get_new_access_token(token_str)

                if refresh_token is None:
                    return response
                elif new_access_token is None:
                    new_access, new_refresh = self.refresh_tokens_directly(
                        refresh_token
                    )
                    # 设置新的token
                    set_new_token(token_str, new_access)
                    set_refresh_token(new_access, new_refresh)
                    new_access_token = new_access

                response["X-Refresh-Token"] = new_access_token
                # 暴露响应
                response["Access-Control-Expose-Headers"] = "X-Refresh-Token"

        except TokenError as e:
            # 打印报错信息
            logger.error("刷新令牌失败", e)
            pass

        return response

    def refresh_tokens_directly(self, refresh_token_str: str):
        """
        直接使用 RefreshToken 类刷新令牌
        Args:
            refresh_token_str: 有效的刷新令牌字符串
        Returns:
            tuple: (新访问令牌, 新刷新令牌)
        """
        try:
            # 验证并创建 RefreshToken 实例
            refresh_token = RefreshToken(token=refresh_token_str)
            refresh_token.verify()
            # 获取新访问令牌
            new_access = str(refresh_token.access_token)

            # 获取用户ID（从令牌中提取）
            user_id = refresh_token[api_settings.USER_ID_CLAIM]
            User = get_user_model()
            user = User.objects.get(id=user_id)

            # 生成新的刷新令牌
            new_refresh = str(RefreshToken.for_user(user))

            # 打印
            print(f"New Access Token: {new_access}")
            print(f"New Refresh Token: {new_refresh}")

            return new_access, new_refresh

        except Exception as e:
            print(f"令牌刷新失败: {e}")
            raise
